#pragma once

#include "LvDevice.hpp"

#include <memory>
#include <unordered_map>
#include <vector>

class LvDescriptorSetLayout
{
public:
    class Builder
    {
    public:
        Builder(LvDevice& lvDevice) : lvDevice {lvDevice} {}

        Builder& addBinding(uint32_t binding,
                            VkDescriptorType descriptorType,
                            VkShaderStageFlags stageFlags,
                            uint32_t count = 1);
        std::unique_ptr<LvDescriptorSetLayout> build() const;

    private:
        LvDevice& lvDevice;
        std::unordered_map<uint32_t, VkDescriptorSetLayoutBinding> bindings {};
    };

    LvDescriptorSetLayout(LvDevice& lvDevice, std::unordered_map<uint32_t, VkDescriptorSetLayoutBinding> bindings);
    ~LvDescriptorSetLayout();
    LvDescriptorSetLayout(const LvDescriptorSetLayout&) = delete;
    LvDescriptorSetLayout& operator=(const LvDescriptorSetLayout&) = delete;

    VkDescriptorSetLayout getDescriptorSetLayout() const { return descriptorSetLayout; }

private:
    LvDevice& lvDevice;
    VkDescriptorSetLayout descriptorSetLayout;
    std::unordered_map<uint32_t, VkDescriptorSetLayoutBinding> bindings;

    friend class LvDescriptorWriter;
};

class LvDescriptorPool
{
public:
    class Builder
    {
    public:
        Builder(LvDevice& lvDevice) : lvDevice {lvDevice} {}

        Builder& addPoolSize(VkDescriptorType descriptorType, uint32_t count);
        Builder& setPoolFlags(VkDescriptorPoolCreateFlags flags);
        Builder& setMaxSets(uint32_t count);
        std::unique_ptr<LvDescriptorPool> build() const;

    private:
        LvDevice& lvDevice;
        std::vector<VkDescriptorPoolSize> poolSizes {};
        uint32_t maxSets                      = 1000;
        VkDescriptorPoolCreateFlags poolFlags = 0;
    };

    LvDescriptorPool(LvDevice& lvDevice,
                      uint32_t maxSets,
                      VkDescriptorPoolCreateFlags poolFlags,
                      const std::vector<VkDescriptorPoolSize>& poolSizes);
    ~LvDescriptorPool();
    LvDescriptorPool(const LvDescriptorPool&) = delete;
    LvDescriptorPool& operator=(const LvDescriptorPool&) = delete;

    bool allocateDescriptor(const VkDescriptorSetLayout descriptorSetLayout, VkDescriptorSet& descriptor) const;

    void freeDescriptors(std::vector<VkDescriptorSet>& descriptors) const;

    void resetPool();

private:
    LvDevice& lvDevice;
    VkDescriptorPool descriptorPool;

    friend class LvDescriptorWriter;
};

class LvDescriptorWriter
{
public:
    LvDescriptorWriter(LvDescriptorSetLayout& setLayout, LvDescriptorPool& pool);

    LvDescriptorWriter& writeBuffer(uint32_t binding, VkDescriptorBufferInfo* bufferInfo);
    LvDescriptorWriter& writeImage(uint32_t binding, VkDescriptorImageInfo* imageInfo);

    bool build(VkDescriptorSet& set);
    void overwrite(VkDescriptorSet& set);

private:
    LvDescriptorSetLayout& setLayout;
    LvDescriptorPool& pool;
    std::vector<VkWriteDescriptorSet> writes;
};
